Write a custom CUDA kernel to optimize `Hard Bootstrapping Loss`.

Formula: Loss = beta * CE(p, target) + (1 - beta) * CE(p, argmax(p))
Where p = softmax(logits).
This is equivalent to: Loss = -beta * log(p_target) - (1 - beta) * log(p_argmax).

Problem Analysis:
1. Memory Intensity: Standard implementation requires calculating Softmax (N, C), finding argmax over the class dimension, and then gathering values. This consumes high memory bandwidth.
2. Redundant Calculation: Finding argmax on logits is equivalent to finding argmax on probabilities. We can fuse this search into the Log-Sum-Exp normalization pass.

Optimization Strategy: Fused Max-ArgMax-Sum Kernel

The goal is to calculate the loss directly from logits without materializing the probability tensor.

1. Mathematical Simplification:
   Using the Log-Sum-Exp trick: log(p_k) = x_k - M - log(S).
   The loss simplifies to: L = log(S) + M - (beta * x_target + (1-beta) * x_argmax).

2. One-Block-per-Row: Each CUDA block handles one sample.

3. Fused Reduction with Index Tracking:
   - Pass 1 (Max & ArgMax): Iterate through the row to find the maximum logit value `M` and its index `idx_max`. This requires a reduction that tracks indices.
   - Pass 2 (Sum): Compute sum of exponentials `S = sum(exp(x - M))`.

4. Vectorized Access: Use `float4` to load data efficiently.

5. Final Computation: Thread 0 combines the results using the simplified formula and writes the scalar loss.
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn
import torch.nn.functional as F

BATCH_SIZE = 2048
NUM_CLASSES = 4096
SHAPE = (BATCH_SIZE, NUM_CLASSES)

BETA = 0.8
REDUCTION = 'none'

class HardBootstrappingLoss(nn.Module):
    """
    Hard Bootstrapping Loss (Reed et al., 2015)
    L = beta * CE(y_true) + (1-beta) * CE(y_pred_max)
    """
    def __init__(self, beta=0.8, reduction='mean'):
        super(HardBootstrappingLoss, self).__init__()
        self.beta = beta
        self.reduction = reduction

    def forward(self, logits: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
        # logits: (N, C)
        # targets: (N)
        ce_loss = F.cross_entropy(logits, targets, reduction='none')
        
        pred_labels = torch.argmax(logits, dim=1)
        
        ce_pred = F.cross_entropy(logits, pred_labels, reduction='none')
        
        loss = self.beta * ce_loss + (1.0 - self.beta) * ce_pred
        
        if self.reduction == 'mean':
            return loss.mean()
        elif self.reduction == 'sum':
            return loss.sum()
        return loss

class Model(nn.Module):
    def __init__(self, beta=0.8, reduction='none'):
        super(Model, self).__init__()
        self.loss_fn = HardBootstrappingLoss(beta=beta, reduction=reduction)
    
    def forward(self, logits, targets):
        return self.loss_fn(logits, targets)

def get_inputs():
    logits = torch.randn(SHAPE, dtype=torch.float32)
    targets = torch.randint(0, NUM_CLASSES, (BATCH_SIZE,), dtype=torch.long)
    return [logits.contiguous(), targets.contiguous()]

def get_init_inputs():
    return [BETA, REDUCTION]